#pragma once
#undef slots  
#include <torch/torch.h>
#include <torch/script.h>
#define slots Q_SLOTS
#include "Vocab.h"

class Seq2SeqModel {
public:
    Seq2SeqModel();
    ~Seq2SeqModel();
    Seq2SeqModel(const std::string& model_path, const Vocab& src_vocab, const Vocab& tar_vocab);
    std::string predict(const std::string& src_sentence, int num_steps = 10);

private:
    torch::jit::script::Module m_model;
    Vocab m_src_vocab;
    Vocab m_tar_vocab;
    torch::Device m_device = torch::cuda::is_available() ? torch::kCUDA : torch::kCPU;

    torch::Tensor preprocess_input(const std::string& src_sentence, int num_steps);
};